import os
from glob import glob
from tqdm import tqdm
from shutil import copy


def get_img_path(img_dir, suffix= ("jpg", "png")):
    img_paths = []
    for sfx in suffix:
        for path in glob(os.path.join(img_dir, f"*.{sfx}")):
            img_paths.append(path)

    return img_paths


def main():
    img_dir = r"D:\lliujian\DeepLearning\datasets\FPC_20250301_ori\images_aug"
    txt_dir = r"D:\lliujian\DeepLearning\datasets\FPC_20250301_ori\images_aug"

    tgt_dir = r"D:\lliujian\DeepLearning\datasets\FPC_20250301_ori\verify"

    save_dir = r"D:\lliujian\DeepLearning\datasets\FPC_20250301_ori\images_aug_select"
    for timg_path in tqdm(get_img_path(tgt_dir)):
        timg_name = os.path.basename(timg_path)
        suffix = os.path.splitext(timg_path)[1]

        img_path = os.path.join(img_dir, timg_name)
        txt_path = os.path.join(txt_dir, timg_name.replace(suffix, ".txt"))
        if os.path.exists(txt_path) and os.path.exists(img_path):
            copy(img_path, os.path.join(save_dir, timg_name))
            copy(txt_path, os.path.join(save_dir, os.path.basename(txt_path)))
        else:
            print(f"{timg_path} not exists!")


if __name__ == '__main__':
    main()